from numpy import shape
import pandas as pd
import operator
from math import log2
from tree_plotter import createPlot

def Ent(data):
    """ 计算集合data的信息熵 公式(4.1) """
    num_sample = len(data)
    label_count = {}
    for sample in data:
        label = sample[-1]
        if label not in label_count.keys():
            label_count[label] = 0
        label_count[label] += 1

    entropy = 0.0
    for key in label_count:
        propotion = float(label_count[key]) / num_sample
        entropy -= propotion * log2(propotion)
    
    return entropy


def split_discrete_data(data, attribute, value):
    """
    将数据集按照离散属性划分
    将数据集中属性为attribute，取值为value的样本划分出来
    data : 待划分数据集
    attribute : 划分所依据的属性
    value : 划分属性取值
    返回的是Dv集合，并且集合中最优划分属性的值被删除
    """
    data_ret = []
    for sample in data:
        if sample[attribute] == value:
            # 在样本中删除该属性
            sample_removed = sample[:attribute]
            sample_removed.extend(sample[attribute+1:])

            data_ret.append(sample_removed)
    
    return data_ret


def split_continuous_data(data, attribute, threshold, direction):
    """
    将数据集按照连续属性划分
    data : 待划分数据集
    attribute : 划分所依据的属性
    threshold : 划分阈值（划分点）
    direction : 划分方向
        0 : 选取低于阈值的样本
        1 : 选取高于阈值的样本
    """
    data_ret = []
    for sample in data:
        if direction == 0:
            if sample[attribute] <= threshold:
                data_ret.append(sample)
        else:
            if sample[attribute] > threshold:
                data_ret.append(sample)
        # 在连续属性上进行划分只是为了评估划分点
        # 从而选择出最优划分点，然后将其转换为离散属性
        # 所以不像“按离散属性划分”那样在划分后将属性删除

    return data_ret


def best_split_feature(data, features):
    """ 选择最优划分属性 """
    num_features = len(data[0]) - 1 # 减去label那一列
    base_ent = Ent(data) # 该集合的信息熵
    best_gain = 0.0 # 记录最大的信息增益
    best_feature = -1 # 记录最优划分属性的索引值
    best_split_dict = {} # 记录连续属性的最佳划分点

    for i in range(num_features):
        # 用第i个属性进行划分，计算其信息增益
        feature_value = [sample[i] for sample in data] # 所有样本第i个属性的所有取值

        if type(feature_value[0]).__name__ == 'float' or type(feature_value[0]).__name__ == 'int':
            # 属性是连续值（P83 4.4.1 连续值处理）
            sorted_value = sorted(feature_value) # 将属性的取值从小到大排序
            dividing_point_list = [] # 候选划分点集合
            for j in range(len(sorted_value) - 1):
                dividing_point_list.append(
                    (sorted_value[j] + sorted_value[j+1]) / 2.0
                )
            
            best_split_ent = float('inf')
            # best_split_ent表示公式(4.8)减号后面的求和项
            # 为了使Gain(D,a,t)最大化，则需要找到使best_split_ent最小的划分点
            for j in range(len(dividing_point_list)):
                point = dividing_point_list[j]
                ent_sum = 0.0
                subdata0 = split_continuous_data(data, i, point, 0)
                subdata1 = split_continuous_data(data, i, point, 1)
                # 公式(4.8)
                weight_negative = len(subdata0) / float(len(data))
                ent_sum += weight_negative * Ent(subdata0)
                weight_positive = len(subdata1) / float(len(data))
                ent_sum += weight_positive * Ent(subdata1)
                if ent_sum < best_split_ent:
                    best_split_ent = ent_sum # 记录当前遇到的最小求和项
                    best_dividing_point_idx = j # 同时记录该划分点的索引值
            
            # 将第i个属性的最佳划分点记录在最佳划分字典中
            best_split_dict[features[i]] = dividing_point_list[best_dividing_point_idx]
            Gain = base_ent - best_split_ent # 计算基于最佳划分点二分后的信息增益（公式(4.8)）
        else:
            # 属性是离散量
            unique_value = set(feature_value) # 第i个属性所有可能的取值（无重复）
            # 下面的算法与属性是连续值的相似（公式(4.2)）
            ent_sum = 0.0
            for value in unique_value:
                subdata = split_discrete_data(data, i, value)
                weight = len(subdata) / float(len(data))
                ent_sum += weight * Ent(subdata)

            Gain = base_ent - ent_sum # 公式(4.2)

        if Gain > best_gain:
            best_feature = i # 记录最优划分属性的索引值
            best_gain = Gain # 记录最优划分属性对应的信息增益

    if type(data[0][best_feature]).__name__ == 'float' or type(data[0][best_feature]).__name__ == 'int':
        # 如果最优划分属性是连续值
        best_split_value = best_split_dict[features[best_feature]] # 读取最优划分属性的最优划分点
        # 把最优划分点的值写到属性里面
        features[best_feature] = features[best_feature] + '<=' + str(best_split_value)
        for i in range(shape(data)[0]):
            # 依据最优划分点，将连续值属性改写为离散值
            # 此时经过处理后按离散值对该集合进行划分
            if data[i][best_feature] <= best_split_value:
                data[i][best_feature] = 1
            else:
                data[i][best_feature] = 0

    return best_feature


def major_class(class_list):
    """ 返回class_list中数量最多的类别 """
    class_count = {}
    for vote in class_list:
        if vote not in class_count.keys():
            class_count[vote] = 0
        class_count[vote] += 1
    
    sorted_class_count = sorted(
        class_count.items(), key=operator.itemgetter(1), reverse=True
    ) # 将类别按票数降序排列
    # 返回票数最多的类别
    return sorted_class_count[0][0]


def TreeGenerate(data, features, data_full, features_full):
    """
    生成决策树
    P74 图 4.2
    data相当于集合D
    features相当于集合A
    """
    class_list = [sample[-1] for sample in data] # 数据的类别列表

    if class_list.count(class_list[0]) == len(class_list):
        # D中样本全属于同一个类别C
        # 标记为C类叶节点
        return class_list[0]

    if len(data[0]) == 1:
        # 待划分属性为空（A为空集），剩下的那一个是label
        # 标记为样本数最多的类
        return major_class(class_list)

    # 选择最优划分属性
    best_feat_idx = best_split_feature(data, features)
    best_feat = features[best_feat_idx]
    my_tree = {best_feat:{}} # 生成分支
    best_feat_value = [sample[best_feat_idx] for sample in data] # 所有样本在最优划分属性上的所有取值
    unique_bf_value = set(best_feat_value) # 最优划分属性上所有可能的取值（无重复）

    if type(data[0][best_feat_idx]).__name__ == 'str':
        # 最优划分属性是离散量
        bf_idx_full = features_full.index(features[best_feat_idx]) # 获取最优划分属性所对应的索引值
        bf_value_full = [sample[bf_idx_full] for sample in data_full] # 所有样本在最优划分属性上的所有取值
        unique_bf_value_full = set(bf_value_full) # 最优划分属性上所有可能的取值（无重复）

    del(features[best_feat_idx]) # 将最优划分属性移除属性集合

    for value in unique_bf_value:
        # 对最优划分属性的每一个取值生成一个分支
        subfeatures = features[:]
        if type(data[0][best_feat_idx]).__name__ == 'str':
            # 最优划分属性是离散量
            unique_bf_value_full.remove(value) # 将最优划分属性的取值从集合中移除
        # 递归生成分支节点
        my_tree[best_feat][value] = TreeGenerate(
            split_discrete_data(data, best_feat_idx, value), subfeatures, data_full, features_full
        )

    if type(data[0][best_feat_idx]).__name__ == 'str':
        # Dv为空集，不进行上面循环中的递归
        # 其类别标记为D中样本最多的类
        for value in unique_bf_value_full:
            my_tree[best_feat][value] = major_class(class_list)
    
    return my_tree


if __name__ == '__main__':
    df = pd.read_csv('watermelon_4_3.csv')
    data = df.values[:, 1:].tolist() # 除编号以外的数据
    data_full = data[:]
    features = df.columns.values[1:-1].tolist() # 样本的属性
    features_full = features[:]

    print('\n决策树')
    my_tree = TreeGenerate(data, features, data_full, features_full)
    print(my_tree)
    createPlot(my_tree, save='信息熵决策树.png')
    